#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
CSV分片工具使用示例
"""

from csv_sharding import CSVSharding
import csv
import os
from pathlib import Path


def create_sample_csv(filename: str = "sample_data.csv", rows: int = 50000):
    """创建一个示例CSV文件用于测试"""
    print(f"创建示例CSV文件: {filename}")
    
    with open(filename, 'w', encoding='utf-8', newline='') as file:
        writer = csv.writer(file)
        
        # 写入表头
        headers = ['ID', '姓名', '邮箱', '电话', '地址', '城市', '备注']
        writer.writerow(headers)
        
        # 写入测试数据
        for i in range(1, rows + 1):
            row = [
                i,
                f'用户{i:05d}',
                f'user{i}@example.com',
                f'138{i%10000:04d}{i%10000:04d}',
                f'测试地址{i}号楼{i%100}室',
                f'城市{i%50}',
                f'这是第{i}行的备注信息，包含一些中文字符和数字{i}'
            ]
            writer.writerow(row)
    
    file_size = os.path.getsize(filename) / (1024 * 1024)
    print(f"生成文件大小: {file_size:.2f} MB")
    return filename


def example_basic_usage():
    """基本使用示例"""
    print("=" * 50)
    print("示例1: 基本使用")
    print("=" * 50)
    
    # 创建示例数据
    sample_file = create_sample_csv("sample_data.csv", 50000)
    
    # 创建分片器（目标大小10MB）
    sharding = CSVSharding(target_size_mb=10.0)
    
    # 执行分片
    output_files = sharding.split_csv(sample_file)
    
    print(f"\n生成的分片文件:")
    for file_path in output_files:
        print(f"  {file_path}")


def example_custom_settings():
    """自定义设置示例"""
    print("\n" + "=" * 50)
    print("示例2: 自定义设置")
    print("=" * 50)
    
    # 创建更大的示例数据
    sample_file = create_sample_csv("large_sample.csv", 100000)
    
    # 创建分片器（目标大小5MB）
    sharding = CSVSharding(target_size_mb=5.0)
    
    # 指定输出目录
    output_dir = "output_shards"
    os.makedirs(output_dir, exist_ok=True)
    
    # 执行分片
    output_files = sharding.split_csv(sample_file, output_dir)
    
    print(f"\n生成的分片文件 (保存在 {output_dir}/):")
    for file_path in output_files:
        print(f"  {Path(file_path).name}")


def example_verify_shards():
    """验证分片结果示例"""
    print("\n" + "=" * 50)
    print("示例3: 验证分片结果")
    print("=" * 50)
    
    # 查找所有分片文件
    shard_files = list(Path(".").glob("*_part*.csv"))
    if not shard_files:
        print("没有找到分片文件，请先运行示例1或2")
        return
    
    print(f"找到 {len(shard_files)} 个分片文件:")
    
    total_rows = 0
    for shard_file in sorted(shard_files):
        with open(shard_file, 'r', encoding='utf-8') as file:
            reader = csv.reader(file)
            rows = sum(1 for _ in reader) - 1  # 减去表头
            size_mb = os.path.getsize(shard_file) / (1024 * 1024)
            
            print(f"  {shard_file.name}: {rows:,} 行, {size_mb:.2f} MB")
            total_rows += rows
    
    print(f"\n总数据行数: {total_rows:,}")


def cleanup_files():
    """清理示例文件"""
    print("\n" + "=" * 50)
    print("清理示例文件")
    print("=" * 50)
    
    # 清理CSV文件
    csv_files = list(Path(".").glob("*.csv"))
    for file in csv_files:
        try:
            file.unlink()
            print(f"删除: {file.name}")
        except Exception as e:
            print(f"删除失败 {file.name}: {e}")
    
    # 清理输出目录
    output_dir = Path("output_shards")
    if output_dir.exists():
        for file in output_dir.glob("*"):
            try:
                file.unlink()
                print(f"删除: {output_dir.name}/{file.name}")
            except Exception as e:
                print(f"删除失败 {file.name}: {e}")
        try:
            output_dir.rmdir()
            print(f"删除目录: {output_dir.name}")
        except Exception as e:
            print(f"删除目录失败: {e}")


if __name__ == "__main__":
    try:
        # 运行示例
        example_basic_usage()
        example_custom_settings()
        example_verify_shards()
        
        # 询问是否清理文件
        print("\n" + "=" * 50)
        response = input("是否清理示例文件? (y/N): ").strip().lower()
        if response in ['y', 'yes', 'Y']:
            cleanup_files()
        else:
            print("保留示例文件")
            
    except KeyboardInterrupt:
        print("\n\n操作被取消")
    except Exception as e:
        print(f"\n错误: {e}")
